#pragma once

#include <fmt/core.h>
#include <set>
#include <cxxabi.h>
#include <filesystem>
#include <fstream>
#include "autograd/variable.h"
#include "log.h"

namespace  ldl
{
namespace autograd
{
class Graph
{
public:
    template<typename T>
    static std::string dot_var(const Variable<T>& v, bool verbose = false) {
        return dot_autograd(*v.m_data_ptr(), verbose);
    }

    template<typename T>
    static std::string dot_autograd(const Autograd<T>& a, bool verbose = false) {
        std::string dot_autograd_str = "{} [label=\"{}\", color=orange, style=filled]\n";
        std::string name = a.name();
        if(verbose) {
            if(!a.name().empty()) {
                name += ": ";
            }
            name += a.str();
        }

        std::string ret = fmt::format(dot_autograd_str, reinterpret_cast<uint64_t>(&a), name);
        return ret;
    }

    static std::string demangle_short_name(const char* mangled_name) {
        int status = 0;
        char* demangled = abi::__cxa_demangle(mangled_name, nullptr, nullptr, &status);
        if (status != 0) {
            return mangled_name; // 解修饰失败，返回原始名称
        }

        std::string result(demangled);
        free(demangled);

        // 移除命名空间
        size_t last_colon = result.find_last_of("::");
        if (last_colon != std::string::npos) {
            result = result.substr(last_colon + 1);
        }

        return result;
    }

    template<typename T>
    static std::string dot_func(const Function<T>& f, bool verbose = false) {
        std::string dot_func = "{} [label=\"{}\", color=lightblue, style=filled, shape=box]\n";
        auto txt = fmt::format(dot_func, reinterpret_cast<uint64_t>(&f), demangle_short_name(typeid(f).name()));
        std::string dot_edge = "{} -> {}\n";
        for(auto& v:f.inputs()) {
            txt += fmt::format(dot_edge, reinterpret_cast<uint64_t>(v.get()), reinterpret_cast<uint64_t>(&f));
        }
        for(auto& v:f.outputs()) {
            txt += fmt::format(dot_edge, reinterpret_cast<uint64_t>(&f), reinterpret_cast<uint64_t>(v.lock().get()));
        }
        return txt;
    }

    template<typename T>
    static std::string get_dot_graph(const Variable<T>& output, bool verbose = true) {
        std::string txt = "";
        std::set<std::shared_ptr<Function<T>>> funcs{};
        std::set<std::shared_ptr<Function<T>>> seen_set{};
        auto add_func = [&](const std::shared_ptr<Function<T>>& f){
            if(seen_set.count(f) <= 0) {
                seen_set.insert(f);
                funcs.insert(f);
            }
        };

        add_func(output.m_data_ptr()->creator());
        txt += dot_var(output,verbose);

        while(!funcs.empty()) {
            auto node = funcs.extract(funcs.begin());
            auto creator = node.value();
            txt += dot_func(*creator);
            for(auto a:creator->inputs()) {
                txt += dot_autograd(*a, verbose);

                if(a->creator() != nullptr) {
                    add_func(a->creator());
                }
            }
        }
        return "digraph g {\n" + txt + "}";
    }

    template<typename T>
    static void plot_dot_graph(const Variable<T>& output, bool verbose = true, const std::string& to_file = "graph.png") {
        auto dot_graph = get_dot_graph(output, verbose);

        const char* home = std::getenv("HOME");
        if (!home) {
            LogInfo() << "Could not get home directory";
            return;
        }

        std::filesystem::path tmp_dir = std::filesystem::path(home) / ".ldl";
        if (!std::filesystem::exists(tmp_dir)) {
            std::filesystem::create_directory(tmp_dir);
        }
        std::filesystem::path graph_path = tmp_dir / "tmp_graph.dot";

        // Write dot graph to file
        std::ofstream out_file(graph_path);
        if (!out_file) {
            LogInfo() << "Could not open file for writing: " << graph_path;
            return;
        }
        out_file << dot_graph;
        out_file.close();

        // Get file extension
        std::filesystem::path to_file_path(to_file);
        std::string extension = to_file_path.extension().string();
        if (!extension.empty()) {
            extension = extension.substr(1); // Remove the dot
        }

        // Execute dot command
        std::string cmd = "dot " + graph_path.string() + " -T " + extension + " -o " + to_file;
        int result = system(cmd.c_str());
        if (result != 0) {
            LogInfo() << "Command failed: " << cmd;
        }
    }
};
}
}
